import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

from globals import *

'''
绘制 三维轨迹图
roll-yaw-pitch
'''
def plot_3d_trajectory(angle_nw, angle_w):
    # 创建图形和轴
    fig = plt.figure()

    pitch_nw = angle_nw[:,1]
    roll_nw = angle_nw[:,0]
    yaw_nw = angle_nw[:,2]

    pitch_w = angle_w[:,1]
    roll_w = angle_w[:,0]
    yaw_w = angle_w[:,2]

    ax = fig.add_subplot(111, projection='3d')
    # 在三维空间中绘制两条曲线
    ax.plot(roll_nw, yaw_nw, pitch_nw,label="nw")
    ax.plot(roll_w, yaw_w, pitch_w,label="w")

    # 设置图形的标签和标题
    ax.set_xlabel('roll')
    ax.set_ylabel('yaw')
    ax.set_zlabel('pitch')
    ax.legend()

    # 显示图形.
    plt.tight_layout()
    plt.show()

'''
绘制 三维气动力图
alpha-beta-C
'''
def plot_3d_wind(input, angle):
    alpha = angle[:,0]
    beta = angle[:,1]

    headers = ["X", "Y", "Z", "Mx", "My", "Mz", "Q"]

    # 创建图形和轴
    fig = plt.figure()

    ax = [None]*6

    for i in range(6):
        ax[i] = fig.add_subplot(2,3,i+1, projection='3d')
        # 在三维空间中绘制两条曲线
        ax[i].plot(alpha, beta, input[:,i])

        # 设置图形的标签和标题
        ax[i].set_xlabel('alpha')
        ax[i].set_ylabel('beta')
        ax[i].set_zlabel(headers[i])
        ax[i].set_title(headers[i])

    # 显示图形.
    plt.tight_layout()
    plt.show()



'''
绘制 时间-气动力 图
'''
def plot_f_time(input, freq):
    input = np.array(input)

    t = np.linspace(0,len(input)*(1/freq),len(input)) 

    headers = ["X", "Y", "Z", "Mx", "My", "Mz", "Q"]

    plt.figure(figsize=FIG_SIZE)
    for i in range(6): # 不画Q
        plt.subplot(6,1,i+1)
        plt.title(headers[i])
        plt.plot(t, input[:,i])

    plt.tight_layout()
    plt.show()


'''
返回 最大周期(最小频率)
'''
def calc_min_f(freq_motion):
    # 将列表转换为集合以获取唯一的元素
    unique_elements = set(freq_motion)

    # 将唯一元素转换为列表并排序
    sorted_unique_elements = sorted(list(unique_elements))

    if(sorted_unique_elements[0]==0):
        f_min = sorted_unique_elements[1]
    else:
        f_min = sorted_unique_elements[0]

    return f_min


'''
绘制 角度-气动力 图
频率：Pitch/Roll/Yaw
陀螺仪：Roll(deg),Pitch(deg),Yaw(deg)
'''
def plot_f_degree(input, angle, freq_motion, freq_angle,avg_flag):
    input = np.array(input)

    t_start = 0
    t_end = len(input)

    if avg_flag == True:
        # 画一个周期
        f_min = calc_min_f(freq_motion)
        print("最小频率："+str(f_min)+"Hz")
        t_start = int(len(input)/2 - 1/f_min * freq_angle/2)
        t_end = int(len(input)/2 + 1/f_min * freq_angle/2)
        input = input[t_start:t_end,:]
        angle = angle[t_start:t_end,:]
        t = t[t_start:t_end]

        print(t_start/freq_angle,t_end/freq_angle)

    pitch = angle[:,1]
    roll = angle[:,0]
    yaw = angle[:,2]

    headers = ["D", "L", "Y", "Mx", "My", "Mz", "Q"]
    plt.figure(figsize=FIG_SIZE)
    if(freq_motion[0]>0):
        for i in range(6):
            plt.subplot(2,3,i+1)
            plt.title(headers[i])
            plt.plot(pitch,input[:,i])
            plt.xlabel('Pitch')
            plt.ylabel('Degree')
    elif(freq_motion[1]>0):
        for i in range(6):
            plt.subplot(2,3,i+1)
            plt.title(headers[i])
            plt.plot(roll,input[:,i])
            plt.xlabel('Roll')
            plt.ylabel('Degree')
    elif(freq_motion[2]>0):
        for i in range(6):
            plt.subplot(2,3,i+1)
            plt.title(headers[i])
            plt.plot(yaw,input[:,i])
            plt.xlabel('Yaw')
            plt.ylabel('Degree')

    plt.tight_layout()
    plt.show()


'''
绘制 气动度-气动力 图
频率：Pitch/Roll/Yaw
陀螺仪：Alpha(deg),Beta(deg),
'''
def plot_f_degree_wind(input, angle, freq_motion, freq_angle,avg_flag):
    input = np.array(input)

    t_start = 0
    t_end = len(input)

    if avg_flag == True:
        # 画一个周期
        f_min = calc_min_f(freq_motion)
        print("最小频率："+str(f_min)+"Hz")
        t_start = int(len(input)/2 - 1/f_min * freq_angle/2)
        t_end = int(len(input)/2 + 1/f_min * freq_angle/2)
        input = input[t_start:t_end,:]
        angle = angle[t_start:t_end,:]

        print(t_start/freq_angle,t_end/freq_angle)

    alpha = angle[:,0]
    beta = angle[:,1]

    headers = ["Drag", "Side", "Lift", "Mx", "My", "Mz", "Q"]
    plt.figure(figsize=FIG_SIZE)

    if(freq_motion[0]>0):
        for i in range(6):
            plt.subplot(2,3,i+1)
            plt.title(headers[i])
            plt.plot(alpha,input[:,i])
            plt.xlabel('Alpha')
            plt.ylabel('Degree')
    else:
        for i in range(6):
            plt.subplot(2,3,i+1)
            plt.title(headers[i])
            plt.plot(beta,input[:,i])
            plt.xlabel('Beta')
            plt.ylabel('Degree')

    plt.tight_layout()
    plt.show()


'''
绘制 气动度-气动力 图
频率：Pitch/Roll/Yaw
陀螺仪：Alpha(deg),Beta(deg),
'''
def plot_f_degree_wind_all(input, angle, freq_motion, freq_angle,avg_flag):
    input = np.array(input)
    
    t = np.linspace(0,len(input)*(1/freq_angle),len(angle)) 

    t_start = 0
    t_end = len(input)

    if avg_flag == True:
        # 画一个周期
        f_min = calc_min_f(freq_motion)
        print("最小频率："+str(f_min)+"Hz")
        t_start = int(len(input)/2 - 1/f_min * freq_angle/2)
        t_end = int(len(input)/2 + 1/f_min * freq_angle/2)
        input = input[t_start:t_end,:]
        angle = angle[t_start:t_end,:]
        t = t[t_start:t_end]

        print(t_start/freq_angle,t_end/freq_angle)

    alpha = angle[:,0]
    beta = angle[:,1]

    headers = ["Drag", "Side", "Lift", "Mx", "My", "Mz", "Q"]

    plt.figure(figsize=FIG_SIZE)

    plt.subplot(3,3,1)
    plt.title("Alpha")
    plt.plot(t,alpha)
    plt.xlabel('time')
    plt.ylabel('Degree')

    plt.subplot(3,3,2)
    plt.title("Beta")
    plt.plot(t,beta)
    plt.xlabel('time')
    plt.ylabel('Degree')

    plt.subplot(3,3,3)
    plt.title("Trajectory")
    plt.plot(beta,alpha)
    plt.xlabel('beta')
    plt.ylabel('alpha')

    if(freq_motion[0]>0):
        for i in range(6):
            plt.subplot(3,3,i+1+3)
            plt.title(headers[i])
            plt.plot(alpha,input[:,i])
            plt.xlabel('Alpha')
            plt.ylabel('Degree')
    else:
        for i in range(6):
            plt.subplot(3,3,i+1+3)
            plt.title(headers[i])
            plt.plot(beta,input[:,i])
            plt.xlabel('Beta')
            plt.ylabel('Degree')

    plt.tight_layout()
    plt.show()

'''
绘制 角度-气动力 动画
'''
def plot_f_degree_ani(input, angle, freq_motion, freq_angle):
    input = np.array(input)

    t_start = 0
    t_end = len(input)

    # 画一个周期
    # f_min = calc_min_f(freq_motion)
    # print("最小频率："+str(f_min)+"Hz")
    # t_start = int(len(input)/2 - 1/f_min * freq_angle/2)
    # t_end = int(len(input)/2 + 1/f_min * freq_angle/2)
    # input = input[t_start:t_end,:]
    # angle = angle[t_start:t_end,:]

    # print(t_start/freq_angle,t_end/freq_angle)

    pitch = angle[:,1]
    roll = angle[:,0]
    yaw = angle[:,2]

    headers = ["D", "L", "Y", "Mx", "My", "Mz", "Q"]

    fig, axs = plt.subplots(2, 3, figsize=(20, 10))

    def update(frame):
        for i in range(6):
            axs[i//3, i%3].clear()
            axs[i//3, i%3].set_title(headers[i])
            if(freq_motion[0]>0):
                axs[i//3, i%3].plot(pitch[:frame], input[:frame, i])
                axs[i//3, i%3].set_xlabel('Pitch')
            elif(freq_motion[1]>0):
                axs[i//3, i%3].plot(roll[:frame], input[:frame, i])
                axs[i//3, i%3].set_xlabel('Roll')
            elif(freq_motion[2]>0):
                axs[i//3, i%3].plot(yaw[:frame], input[:frame, i])
                axs[i//3, i%3].set_xlabel('Yaw')
            axs[i//3, i%3].set_ylabel('Degree')

    ani = FuncAnimation(fig, update, frames=range(0, len(input), 30), interval=100)
    # ani.save('animation.gif', writer='pillow')

    plt.tight_layout()
    plt.show()


'''
绘制 时间-绳拉力 图
'''
def plot_tension(input, freq):
    input = np.array(input)

    t = np.linspace(0,len(input)*freq,len(input)) 

    plt.figure(figsize=FIG_SIZE)
    for i in range(8):
        plt.subplot(len(input[0]),1,i+1)
        plt.plot(t, input[:,i])

    plt.tight_layout()
    plt.show()


'''
绘制 时间-编码器 图
'''
def plot_encoder(input, freq):
    input = np.array(input)

    t = np.linspace(0,len(input)*(1/freq),len(input)) 

    plt.figure(figsize=FIG_SIZE)
    for i in range(17):
        plt.subplot(len(input[0]),1,i+1)
        plt.plot(t, input[:,i])

    plt.tight_layout()
    plt.show()


'''
绘制 时间-陀螺仪 图
'''
def plot_angle(input, freq):
    input = np.array(input)

    t = np.linspace(0,len(input)*(1/freq),len(input)) 

    headers = ["Roll", "Pitch", "Yaw"]

    plt.figure(figsize=FIG_SIZE)
    for i in range(3):
        plt.subplot(3,1,i+1)
        plt.title(headers[i])
        plt.plot(t, input[:,i])

    plt.tight_layout()
    plt.show()


'''
绘制 有风无风 角度对比
'''
def plot_comp(nw,w,freq):
    w = np.array(w)
    nw = np.array(nw) 

    t_w = np.linspace(0,len(w)*(1/freq),len(w)) 
    t_nw = np.linspace(0,len(nw)*(1/freq),len(nw)) 

    headers = ["Roll", "Pitch", "Yaw"]

    plt.figure(figsize=FIG_SIZE)
    for i in range(3):
        plt.subplot(3,1,i+1)
        plt.title(headers[i])
        plt.plot(t_nw, nw[:,i])
        plt.plot(t_w, w[:,i])

    plt.tight_layout()
    plt.show()


'''
绘制 全部 时间图
有风无风角度/滤波后气动力/角度-气动力
标题写上运动频率
'''
def plot_all_time(angle_nw,angle_w,freq_angle,F_n_Q_body,F_n_Q_wind,freq_aero):
    
    t_w = np.linspace(0,len(angle_w)*(1/freq_angle),len(angle_w)) 
    t_nw = np.linspace(0,len(angle_nw)*(1/freq_angle),len(angle_nw)) 

    headers = ["Roll", "Pitch", "Yaw"]

    plt.figure(figsize=FIG_SIZE)
    for i in range(3):
        plt.subplot(3,3,i+1)
        plt.title(headers[i])
        plt.plot(t_nw, angle_nw[:,i],'--',label="nw",color="r")
        plt.plot(t_w, angle_w[:,i],label='w',color='blue')
        plt.legend()
    

    F_n_Q_wind = np.array(F_n_Q_wind)
    t = np.linspace(0,len(F_n_Q_wind)*(1/freq_aero),len(F_n_Q_wind)) 

    headers = ["X", "Y", "Z", "Mx", "My", "Mz", "Q"]

    for i in range(6): # 不画Q
        plt.subplot(3,3,i+4)
        plt.title(headers[i])
        plt.plot(t, F_n_Q_body[:,i],'--',label="body",color="r")
        plt.plot(t, F_n_Q_wind[:,i],label="wind",color='blue')
        plt.legend()

    plt.tight_layout()
    plt.show()


'''
绘制 全部 角度 图
'''
def plot_all_angle(angle_nw,angle_w,freq_angle,F_n_Q_body,F_n_Q_wind,freq_aero,freq_motion,angle):
    
    t_w = np.linspace(0,len(angle_w)*(1/freq_angle),len(angle_w)) 
    t_nw = np.linspace(0,len(angle_nw)*(1/freq_angle),len(angle_nw)) 

    headers = ["Roll", "Pitch", "Yaw"]

    plt.figure(figsize=FIG_SIZE)
    for i in range(3):
        plt.subplot(3,3,i+1)
        plt.title(headers[i])
        plt.plot(t_nw, angle_nw[:,i],'--',label="nw",color="r")
        plt.plot(t_w, angle_w[:,i],label='w',color='blue')
        plt.legend()
    

    F_n_Q_wind = np.array(F_n_Q_wind)
    t = np.linspace(0,len(F_n_Q_wind)*(1/freq_aero),len(F_n_Q_wind)) 

    pitch = angle[:,1]
    roll = angle[:,0]
    yaw = angle[:,2]

    headers = ["X", "Y", "Z", "Mx", "My", "Mz", "Q"]

    if(freq_motion[0]>0):
        for i in range(6):
            plt.subplot(3,3,i+1+3)
            plt.title(headers[i])
            plt.plot(pitch,F_n_Q_body[:,i])
            plt.xlabel('Pitch')
            plt.ylabel('Degree')
    elif(freq_motion[1]>0):
        for i in range(6):
            plt.subplot(3,3,i+1+3)
            plt.title(headers[i])
            plt.plot(roll,F_n_Q_body[:,i])
            plt.xlabel('Roll')
            plt.ylabel('Degree')
    elif(freq_motion[2]>0):
        for i in range(6):
            plt.subplot(3,3,i+1+3)
            plt.title(headers[i])
            plt.plot(yaw,F_n_Q_body[:,i])
            plt.xlabel('Yaw')
            plt.ylabel('Degree')

    plt.tight_layout()
    plt.show()

'''
绘制 全部 角度 动图
'''
def plot_all_angle_ani(angle_nw, angle_w, freq_angle, F_n_Q_body, F_n_Q_wind, freq_aero, freq_motion, angle):
    t_w = np.linspace(0, len(angle_w) * (1 / freq_angle), len(angle_w))
    t_nw = np.linspace(0, len(angle_nw) * (1 / freq_angle), len(angle_nw))

    headers = ["Roll", "Pitch", "Yaw"]

    fig, axs = plt.subplots(3, 3, figsize=(20, 10))

    headers_aero = ["X", "Y", "Z", "Mx", "My", "Mz", "Q"]

    F_n_Q_wind = np.array(F_n_Q_wind)
    t = np.linspace(0, len(F_n_Q_wind) * (1 / freq_aero), len(F_n_Q_wind))

    pitch = angle[:,1]
    roll = angle[:,0]
    yaw = angle[:,2]



    def update(frame):
        for i in range(3):
            axs[i, 0].clear()
            axs[i, 0].set_title(headers[i])
            axs[i, 0].plot(t_nw[:frame], angle_nw[:frame, i], '--', label="nw", color="r")
            axs[i, 0].plot(t_w[:frame], angle_w[:frame, i], label='w', color='blue')
            axs[i, 0].legend()

        # for i in range(3,9):
        #     axs[i//3, i%3].clear()
        #     axs[i//3, i%3].set_title(headers_aero[i])
        #     if(freq_motion[0]>0):
        #         axs[i//3, i%3].plot(pitch[:frame], F_n_Q_body[:frame, i])
        #         axs[i//3, i%3].set_xlabel('Pitch')
        #     elif(freq_motion[1]>0):
        #         axs[i//3, i%3].plot(roll[:frame], F_n_Q_body[:frame, i])
        #         axs[i//3, i%3].set_xlabel('Roll')
        #     elif(freq_motion[2]>0):
        #         axs[i//3, i%3].plot(yaw[:frame], F_n_Q_body[:frame, i])
        #         axs[i//3, i%3].set_xlabel('Yaw')
        #     axs[i//3, i%3].set_ylabel('Degree')

        plt.tight_layout()

    ani = FuncAnimation(fig, update, frames=len(angle), interval=50)
    plt.tight_layout()
    plt.show()



'''
绘制 滤波后的电压
'''
def plot_u(u_nw_raw, u_w_raw, u_nw, u_w, freq):
    u_nw_raw = np.array(u_nw_raw)
    u_w_raw = np.array(u_w_raw)
    u_nw = np.array(u_nw)
    u_w = np.array(u_w)

    t = np.linspace(0,len(u_w)*(1/freq),len(u_w)) 

    headers = ["u1", "u2", "u3", "u4", "u5", "u6"]

    plt.figure(figsize=FIG_SIZE)
    for i in range(6): # 不画Q
        plt.subplot(6,2,2*i+1)
        plt.title(headers[i])
        plt.plot(t, u_nw_raw[:,i])
        plt.plot(t, u_nw[:,i])

    for i in range(6): # 不画Q
        plt.subplot(6,2,2*i+2)
        plt.title(headers[i])
        plt.plot(t, u_w_raw[:,i])
        plt.plot(t, u_w[:,i])
    plt.tight_layout()
    plt.show()

'''
绘制 全部 图
有风无风角度/滤波后气动力/角度-气动力

标题写上运动频率
'''
def plot_all(angle_nw,angle_w,freq_angle,F_n_Q_body,F_n_Q_wind,freq_aero,freq_motion,angle):
    
    t_w = np.linspace(0,len(angle_w)*(1/freq_angle),len(angle_w)) 
    t_nw = np.linspace(0,len(angle_nw)*(1/freq_angle),len(angle_nw)) 

    headers = ["Roll", "Pitch", "Yaw"]

    plt.figure(figsize=FIG_SIZE)
    for i in range(3):
        plt.subplot(5,3,i+1)
        plt.title(headers[i])
        plt.plot(t_nw, angle_nw[:,i],'--',label="nw",color="r")
        plt.plot(t_w, angle_w[:,i],label='w',color='blue')
        plt.legend()
    

    F_n_Q_wind = np.array(F_n_Q_wind)
    t = np.linspace(0,len(F_n_Q_wind)*(1/freq_aero),len(F_n_Q_wind)) 

    headers = ["X", "Y", "Z", "Mx", "My", "Mz", "Q"]

    for i in range(6): # 不画Q
        plt.subplot(5,3,i+4)
        plt.title(headers[i])
        plt.plot(t, F_n_Q_body[:,i],'--',label="body",color="r")
        plt.plot(t, F_n_Q_wind[:,i],label="wind",color='blue')
        plt.legend()

    pitch = angle[:,1]
    roll = angle[:,0]
    yaw = angle[:,2]

    if(freq_motion[0]>0):
        for i in range(6):
            plt.subplot(5,6,i+1+18)
            plt.title(headers[i])
            plt.plot(pitch,F_n_Q_body[:,i])
            plt.xlabel('Pitch')
            plt.ylabel('Degree')
    elif(freq_motion[1]>0):
        for i in range(6):
            plt.subplot(5,6,i+1+18)
            plt.title(headers[i])
            plt.plot(roll,F_n_Q_body[:,i])
            plt.xlabel('Roll')
            plt.ylabel('Degree')

    elif(freq_motion[2]>0):
        for i in range(6):
            plt.subplot(5,6,i+1+18)
            plt.title(headers[i])
            plt.plot(yaw,F_n_Q_wind[:,i])
            plt.xlabel('Yaw')
            plt.ylabel('Degree')

    if(freq_motion[0]>0):
        for i in range(6):
            plt.subplot(5,6,i+1+24)
            plt.title(headers[i])
            plt.plot(pitch,F_n_Q_wind[:,i])
            plt.xlabel('Pitch')
            plt.ylabel('Degree')
    elif(freq_motion[1]>0):
        for i in range(6):
            plt.subplot(5,6,i+1+24)
            plt.title(headers[i])
            plt.plot(roll,F_n_Q_wind[:,i])
            plt.xlabel('Roll')
            plt.ylabel('Degree')

    elif(freq_motion[2]>0):
        for i in range(6):
            plt.subplot(5,6,i+1+24)
            plt.title(headers[i])
            plt.plot(yaw,F_n_Q_wind[:,i])
            plt.xlabel('Yaw')
            plt.ylabel('Degree')
    plt.tight_layout()
    plt.show()